# -*- coding: utf-8 -*-
"""
数据导出工具模块

此模块提供用于导出计算结果的函数。
主要功能包括：
- 导出数据到CSV文件
- 导出数据到NetCDF文件
- 导出数据到Excel文件
"""

import os
import pandas as pd
import xarray as xr
import numpy as np
from pathlib import Path


def export_to_csv(data, output_file, index=True, header=True):
    """
    导出数据到CSV文件
    
    参数:
        data (pandas.DataFrame): 要导出的数据
        output_file (str): 输出文件路径
        index (bool): 是否包含索引
        header (bool): 是否包含表头
        
    返回:
        bool: 是否成功导出
    """
    try:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        
        # 导出数据到CSV文件
        data.to_csv(output_file, index=index, header=header)
        
        print(f"数据已导出到: {output_file}")
        return True
    except Exception as e:
        print(f"导出数据到CSV文件失败: {e}")
        return False


def export_to_excel(data, output_file, sheet_name='Sheet1', index=True, header=True):
    """
    导出数据到Excel文件
    
    参数:
        data (pandas.DataFrame): 要导出的数据
        output_file (str): 输出文件路径
        sheet_name (str): 工作表名称
        index (bool): 是否包含索引
        header (bool): 是否包含表头
        
    返回:
        bool: 是否成功导出
    """
    try:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        
        # 导出数据到Excel文件
        data.to_excel(output_file, sheet_name=sheet_name, index=index, header=header)
        
        print(f"数据已导出到: {output_file}")
        return True
    except Exception as e:
        print(f"导出数据到Excel文件失败: {e}")
        return False


def export_to_netcdf(data, output_file, variable_name='data'):
    """
    导出数据到NetCDF文件
    
    参数:
        data (xarray.Dataset or xarray.DataArray): 要导出的数据
        output_file (str): 输出文件路径
        variable_name (str): 变量名称，仅当data为numpy.ndarray时使用
        
    返回:
        bool: 是否成功导出
    """
    try:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        
        # 如果数据是numpy.ndarray，则转换为xarray.DataArray
        if isinstance(data, np.ndarray):
            # 创建坐标
            coords = {}
            for i, dim in enumerate(data.shape):
                coords[f'dim_{i}'] = np.arange(dim)
            
            # 创建DataArray
            data = xr.DataArray(data, coords=coords, name=variable_name)
        
        # 导出数据到NetCDF文件
        if isinstance(data, xr.DataArray):
            data.to_netcdf(output_file)
        else:
            data.to_netcdf(output_file)
        
        print(f"数据已导出到: {output_file}")
        return True
    except Exception as e:
        print(f"导出数据到NetCDF文件失败: {e}")
        return False


def export_grid_points(grid_points, output_file, province_name=None):
    """
    导出网格点数据到CSV文件
    
    参数:
        grid_points (list): 网格点列表，每个元素为(lon, lat)元组
        output_file (str): 输出文件路径
        province_name (str): 省份名称，如为None则不包含省份列
        
    返回:
        bool: 是否成功导出
    """
    try:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        
        # 创建DataFrame
        df = pd.DataFrame(grid_points, columns=['longitude', 'latitude'])
        
        # 如果指定了省份名称，则添加省份列
        if province_name is not None:
            df['province'] = province_name
        
        # 导出数据到CSV文件
        df.to_csv(output_file, index=False)
        
        print(f"网格点数据已导出到: {output_file}")
        return True
    except Exception as e:
        print(f"导出网格点数据失败: {e}")
        return False


def export_drought_indices(indices, output_file, index_type, region=None, grid_point=None):
    """
    导出干旱指数数据到CSV文件
    
    参数:
        indices (dict): 干旱指数字典，格式为{year: {month: value}}
        output_file (str): 输出文件路径
        index_type (str): 指数类型，如'DWAAI'、'LDFAI'等
        region (str): 区域名称，如为None则不包含区域列
        grid_point (tuple): 网格点坐标(lon, lat)，如为None则不包含网格点列
        
    返回:
        bool: 是否成功导出
    """
    try:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        
        # 创建数据列表
        data_list = []
        
        # 遍历年份和月份
        for year, year_data in indices.items():
            for month, value in year_data.items():
                # 创建数据项
                data_item = {
                    'year': year,
                    'month': month,
                    index_type: value
                }
                
                # 如果指定了区域名称，则添加区域列
                if region is not None:
                    data_item['region'] = region
                
                # 如果指定了网格点坐标，则添加网格点列
                if grid_point is not None:
                    data_item['longitude'] = grid_point[0]
                    data_item['latitude'] = grid_point[1]
                
                # 添加到数据列表
                data_list.append(data_item)
        
        # 创建DataFrame
        df = pd.DataFrame(data_list)
        
        # 导出数据到CSV文件
        df.to_csv(output_file, index=False)
        
        print(f"干旱指数数据已导出到: {output_file}")
        return True
    except Exception as e:
        print(f"导出干旱指数数据失败: {e}")
        return False